import numpy as np
import torch
from torch.utils.data import TensorDataset,DataLoader
import matplotlib.pylab as plt

class  TorchModel(torch.nn.Module):
    def __init__(self):
        super(TorchModel, self).__init__()

        self.conv1 = torch.nn.Conv2d(1,6,5,1) # 输入通道，输出通道，卷积， 步长
        self.pool1 = torch.nn.MaxPool2d(2,2)

        self.conv2 = torch.nn.Conv2d(6,12,3,1)
        self.pool2 = torch.nn.MaxPool2d(2,2)

        self.fc3 = torch.nn.Linear(2*2*12,64)
        self.droup1 = torch.nn.Dropout(0.3)

        self.fc4 = torch.nn.Linear(64,10)

    def forward(self,inputs):

        out = torch.nn.ReLU()(self.conv1(inputs))
        out = self.pool1(out)
        out = torch.nn.ReLU()(self.conv2(out))
        out = self.pool2(out)
        out = out.view(out.size(0),-1)
        out = torch.nn.ReLU()(self.fc3(out))
        out = self.droup1(out)
        out = self.fc4(out)

        return out

    def loss(self,h,y):

        return torch.nn.CrossEntropyLoss()(h,y)

    # 注解的含义是将下面的函数作为一个属性进行处理
    @property
    def optimizer(self):
        return torch.optim.SGD(self.parameters(),lr=0.001)

    def metric(self,h,y):

        pridction = torch.max(h,1)[1]
        acc = (pridction==y).float().mean()
        return acc
if __name__ == '__main__':

    # 1.将数据集进行读取

    data = np.loadtxt('../img_16_10k.txt',delimiter=',')
    # 2.数据集进行切分
    x = data[:,:-1]
    y = data[:,-1]

    x = torch.tensor(x,dtype=torch.float32)
    y = torch.tensor(y,dtype=torch.long)


    x = torch.reshape(x,shape=(-1,1,16,16))

    print(x.shape)
    print(y.shape)

    data = TensorDataset(x,y)
    train_data = DataLoader(data,shuffle=True,batch_size=100)

    # x_,y_ = next(iter(train_data))
    # 3.分别使用pytorch和keras进行卷积处理
    model = TorchModel()

    list = []
    for i in range(20):
        loss_list = []
        for img,label in train_data: # 小批量
            img_ = img
            label_ = label
            h = model(img_) # 训练集的预测值
            loss_val = model.loss(h,label_)  # 当前的训练损失会
            loss_val.backward()
            model.optimizer.step()
            model.optimizer.zero_grad()
            loss_list.append(loss_val.item())
        avg_loss = np.mean(loss_list) # 验证预测取均值
        list.append(avg_loss) # 添加到数据中
        per_y = model(x)
        acc = model.metric(per_y,y) #
        print(acc,avg_loss)
    plt.plot(list,'--or')
    plt.show()


